import pathlib
from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess

model_dir = "iic/SenseVoiceSmall"

model = AutoModel(
    model=model_dir,
    trust_remote_code=True,
    remote_code="./model.py",
    vad_model="fsmn-vad",
    vad_kwargs={"max_single_segment_time": 30000},
    device="cuda:0",
)

# zh
res = model.generate(
    input=f"{model.model_path}/example/zh.mp3",
    cache={},
    language="auto",  # "zh", "en", "yue", "ja", "ko", "nospeech"
    use_itn=True,
    batch_size_s=60,
    merge_vad=True,  #
    merge_length_s=15,
)
text = rich_transcription_postprocess(res[0]["text"])
print(text)

# 本地模型目录
# model_dir = str(pathlib.Path.home()) + "/.cache/modelscope/hub/iic"

# model = AutoModel(
#     model=model_dir + "/SenseVoiceSmall",
#     trust_remote_code=True,
#     disable_update=True,
#     remote_code="./SenseVoice/model.py",
#     vad_model=model_dir + "/speech_fsmn_vad_zh-cn-16k-common-pytorch",
#     vad_kwargs={"max_single_segment_time": 30000},
#     device="cuda:0",
# )


# res = model.generate(
#     input=model_dir + "/SenseVoiceSmall/example/zh.mp3",
#     cache={},
#     language="auto", 
#     use_itn=True,
#     batch_size_s=60,
#     merge_vad=True, 
#     merge_length_s=15,
# )
# text = rich_transcription_postprocess(res[0]["text"])
# print(text)